package com.sicheng.lc.周赛.分类.堆;

/**
 * @author zsc
 * @version 1.0
 * @date 2022/6/28 12:39
 */
public class 向数组中追加K个整数 {
    //https://leetcode.cn/problems/append-k-integers-with-minimal-sum/
    public long minimalKSum(int[] nums, int k) {
        heapSorted(nums);
        long res = 0;
        if (nums[0] > 1) {
            res += (long) (1 + Math.min(nums[0] - 1, k)) * Math.min(k, nums[0] - 1) >> 1;
            k -= Math.min(k, nums[0] - 1);
        }
        if (k > 0)
            for (int i = 0; i < nums.length; i++) {
                if (i + 1 < nums.length && nums[i + 1] > nums[i]) {
                    int m = nums[i + 1] - nums[i] - 1;
                    res += (long) (nums[i] + 1 + Math.min(nums[i + 1] - 1, nums[i] + k)) * Math.min(k, m) >> 1;
                    k -= Math.min(k, m);
                }
                if (k == 0)
                    break;
            }

        if (k > 0)
            res += (long) (nums[nums.length - 1] + 1 + nums[nums.length - 1] + k) * k >> 1;
        return res;
    }

    void heapSorted(int[] nums) {
        int n = nums.length;
        heapy(nums, nums.length);
        while (n-- > 0) {
            swap(nums, n);
            shiftDown(nums, n, nums[0], 0);
        }
    }

    private void swap(int[] nums, int j) {
        int t = nums[0];
        nums[0] = nums[j];
        nums[j] = t;
    }

    void heapy(int[] nums, int size) {
        for (int i = (size >> 1) - 1; i >= 0; i--) {
            shiftDown(nums, size, nums[i], i);
        }
    }

    void shiftDown(int[] nums, int size, int e, int k) {
        int half = size >> 1;
        while (k < half) {
            int c = k << 1 | 1;
            if (c + 1 < size && nums[c + 1] > nums[c])
                c++;

            if (nums[c] <= e) {
                break;
            }
            nums[k] = nums[c];
            k = c;
        }
        nums[k] = e;
    }

    public static void main(String[] args) {
        向数组中追加K个整数 s = new 向数组中追加K个整数();
        int[] nums = {1, 4, 25, 10, 25};

        System.out.println(s.minimalKSum(nums, 2));
    }
}
